package cn.iocoder.yudao.module.infra.embedding.impl.dashscope;

import cn.hutool.core.util.StrUtil;
import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
import cn.iocoder.yudao.module.infra.embedding.EmbeddingConfig;
import cn.iocoder.yudao.module.infra.embedding.IEmbeddingService;
import cn.iocoder.yudao.module.infra.okhttp.OkHttpUtils;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import okhttp3.MediaType;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;
import org.springframework.stereotype.Service;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * https://bailian.console.aliyun.com/?tab=model#/api-key
 * https://bailian.console.aliyun.com/?apiKey=1
 * https://help.aliyun.com/zh/model-studio/developer-reference/text-embedding-synchronous-api?spm=a2c4g.11186623.help-menu-2400256.d_2_5_0.35877980bznekD
 */
@Slf4j
@Service("embedding-dashscope")
public class DashScopeEmbeddingServiceImpl implements IEmbeddingService {
    private EmbeddingConfig.Config dashscopeConfig;

    @Resource
    private EmbeddingConfig embeddingConfig;

    @PostConstruct
    public void init(){
        this.dashscopeConfig = embeddingConfig.getServices().get("dashscope");
    }

    @Override
    public List<? extends Number> create(String text) {
        final long startTime = System.currentTimeMillis();
        final String apiKey = dashscopeConfig.getApiKey();
        String apiUrl = dashscopeConfig.getApiUrl();

        final Map<String, Object> req = new HashMap<>();
        req.put("model", dashscopeConfig.getModel());
        req.put("input", text);
        req.put("dimensions", dashscopeConfig.getDimensions());
        req.put("encoding_format", "float");
        final String body = JsonUtils.toJsonString(req);

        final RequestBody requestBody = RequestBody.create(
                MediaType.parse("application/json"),
                body
        );

        final Request request = new Request.Builder()
                .url(apiUrl)
                .header("Content-Type", "application/json")
                .header("Authorization", "Bearer " + apiKey)
                .post(requestBody)
                .build();
        log.info("请求Embedding-开始, apiUrl = {}, body = {}", apiUrl, body);

        try (Response response = OkHttpUtils.getEmbeddingClient().newCall(request).execute()) {
            final String responseBody = response.body() != null ? response.body().string() : null;
            if (response.isSuccessful() && response.body() != null) {
                final long duration = System.currentTimeMillis() - startTime;
                log.info("请求Embedding-成功, duration = {}", duration);
                final DashScopeEmbeddingRes res = JsonUtils.parseObject(responseBody, DashScopeEmbeddingRes.class);
                return res.getData().get(0).getEmbedding();
            } else {
                log.error("请求Embedding-失败, code = {}, body = {}", response.code(), responseBody);
                return null;
            }
        } catch (Exception e) {
            log.error(StrUtil.format("请求Embedding-异常, apiUrl = {}, body = {}", apiUrl, body), e);
            return null;
        }
    }


    public static void main(String[] args) {
        // 生成DashscopeEmbeddingServiceImpl的测试代码
        EmbeddingConfig.Config config = new EmbeddingConfig.Config();
        config.setDimensions(768);
        config.setModel("text-embedding-v3");
        config.setApiKey("sk-e97333b592b84ac18b97688759094454");
        config.setApiUrl("https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings");

        DashScopeEmbeddingServiceImpl service = new DashScopeEmbeddingServiceImpl();
        service.embeddingConfig = new EmbeddingConfig();
        service.embeddingConfig.setPrimary("dashscope");
        service.embeddingConfig.setServices(new HashMap<>());
        service.embeddingConfig.getServices().put("dashscope", config);

        List<? extends Number> res = service.create("你好");
        System.out.println(res);
    }

}
